from elasticsearch import Elasticsearch
from elasticsearch import helpers
import json

import requests

from parser_file import load_pdf,  text_splitter
def search_embedding(text):
    data = {
        "query": text
    }
    url = "http://ip:port/user"

    response = requests.post(url, json=data)
    result = json.loads(response.text)
    return result


class ElasticCreateModule:
    """
    基于es构建索引，可以支持增加、删除功能；
    """
    def __init__(self, corpus_id=""):
        self.es = Elasticsearch(hosts=["http://user:pass@user:port/"])
        self.text_analyzer = "standard"
        self.text_search_analyzer = "standard"
        self.es_index_name = "rfp-project"

        self.embedding_dim = len(search_embedding("text"))
        print(self.embedding_dim)

    def build_index(self):
        """
        :param index_name: 索引名称
        :param delete_old_index:  是否删除已存在同名索引
        :return:
        """

        delete_old_index = True

        self._index_mappings = {
            "mappings": {
                "properties": {
                    "text-vector": {
                        "type": "dense_vector",
                        "dims": self.embedding_dim,
                        "index": True,
                        "similarity": "cosine"
                    },
                    "text": {
                        "type": "text",
                        "index": True,
                        "analyzer": self.text_analyzer,
                        "search_analyzer": self.text_search_analyzer,
                    },
                    "docName": {
                        "type": "text",
                        "index": True
                    }
                }
            }
        }

        index_name = self.es_index_name.strip()

        print("exists:", self.es.indices.exists(index=index_name))
        if not self.es.indices.exists(index=index_name):
            self.es.indices.create(index=index_name, body=self._index_mappings)
            print("create index: {}".format(index_name))
        else:
            if delete_old_index:
                self.es.indices.delete(index=index_name)
                res = self.es.indices.create(index=index_name, body=self._index_mappings)

    def add_datas(self, datas):
        """
        [[doc_id, [[para_id,context]]]]
        :param index_name: 索引名称
        :param batch_data: 需索引数据  类似 {"paragraph": obj["title"] + ";" + obj["text"]} paragraph需要和index对应
        :return:
        """
        batch_data = []
        for doc_path in datas:
            doc_content = load_pdf(doc_path)
            doc_name = doc_path.split("/")[-1]
            print(len(doc_content))
            contents =[_.page_content for _ in doc_content]
            for con in contents:
                embedding = search_embedding(con)
                batch_data.append({
                    "text": con,
                    "text-vector": embedding,
                    "docName": doc_name
                })
            # break

        index_name = self.es_index_name
        index_data = []
        for _data in batch_data:

            body = {"_index": index_name, "_source": _data}
            index_data.append(body)

        var_data = []
        for _ in index_data:
            var_data.append(_)
            if len(var_data) == 10:
                print(json.dumps(var_data))
                helpers.bulk(self.es, actions=var_data)
                var_data = []

        res = helpers.bulk(self.es, actions=var_data)
        print(res)
        self.es.indices.refresh()

    # def search_by_text(self, query, doc_num=50, query_key="text"):
    #     """
    #     :param doc_num: 返回候选个数
    #     :param index_name: 索引名称
    #     :param query: 搜索query
    #     :return:
    #     """
    #
    #     index_name = self.es_index_name
    #     print("index_name",  index_name)
    #     res = self.es.search(index=index_name, body={"query": {"match": {query_key: query}}}, _source=["text", "docId", "docTitle", "paraId"], size=doc_num)
    #     for i, hit  in enumerate(res["hits"]["hits"]):
    #         print("第 {}条： {}".format(i, hit["_source"]["text"]))
    #     return res

import os
if __name__=="__main__":

    #创建索引
    es_client = ElasticCreateModule()
    filep = ["../DATA/rule_process/{}".format(filename) for filename in os.listdir("../DATA/rule_process/") if filename.endswith("pdf")]
    es_client.build_index()
    es_client.add_datas(filep)
